# !/usr/bin/env python
# -*- coding:utf-8 -*-
"""
Created on Tue Nov 5 11:56:06 2020
@author: hongwang (hongwang01@stu.xjtu.edu.cn)
MICCAI2021: ``InDuDoNet: An Interpretable Dual Domain Network for CT Metal Artifact Reduction''
paper link： https://arxiv.org/pdf/2109.05298.pdf
"""
from __future__ import print_function

import argparse

from network.build_gemotry import generate_rat_femur_odl_operator
from utils.dataloader import RatFemurInDuDoNetMultiInferDataLoader
import os
import time
from math import ceil

import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import torch.nn.parallel
import torch.optim as optim
import yaml
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader

from utils.EvaluationMetrics import computeMAE, computePSNR, computeSSIM
from utils.utils import InferenceResultSaver
from network.InDuDoNet import InDuDoNet
from network.InDuDoNetRatFemur import InDuDoNetRatFemur
from utils.FBP import AstraReconstruction
from utils.Visualizer import Visualizer
from utils.config import Config
from utils.dataloader import GatedConvMultiDataLoader1, RatFemurInDuDoNetMultiDataLoader

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
parser = argparse.ArgumentParser()
parser.add_argument("--data_path", type=str, default=r"D:\gs\code\InDuDoNet\deeplesion\train",
                    help='txt path to training spa-data')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=0)
parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
parser.add_argument('--patchSize', type=int, default=364, help='the height / width of the input image to network')
parser.add_argument('--niter', type=int, default=100, help='total number of training epochs')
parser.add_argument('--batchnum', type=int, default=1000,
                    help='batchsize*batchnum=1000 for randomly selecting 1000 imag pairs at every iteration')
parser.add_argument('--num_channel', type=int, default=32,
                    help='the number of dual channels')  # refer to https://github.com/hongwang01/RCDNet for the channel concatenation strategy
parser.add_argument('--T', type=int, default=4, help='the number of ResBlocks in every ProxNet')
parser.add_argument('--S', type=int, default=10, help='the number of total iterative stages')
parser.add_argument('--resume', type=int, default=0, help='continue to train')
parser.add_argument("--milestone", type=int, default=[40, 80], help="When to decay learning rate")
parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate')
parser.add_argument('--log_dir', default='./logs/', help='tensorboard logs')
parser.add_argument('--model_dir', default='./models/', help='saving model')
parser.add_argument('--eta1', type=float, default=1, help='initialization for stepsize eta1')
parser.add_argument('--eta2', type=float, default=5, help='initialization for stepsize eta2')
parser.add_argument('--alpha', type=float, default=0.5, help='initialization for weight factor')
parser.add_argument('--gamma', type=float, default=1e-1, help='hyper-parameter for balancing different loss items')
opt = parser.parse_args()

# create path
try:
    os.makedirs(opt.model_dir)
except OSError:
    pass

cudnn.benchmark = True


def inference():
    op_module_fp, op_module_pT = generate_rat_femur_odl_operator()
    config_file_path = r"../Yaml/RatFemurInDuDoNetInference.yaml"
    config = Config(config_file_path)
    inference_dataloader = RatFemurInDuDoNetMultiInferDataLoader(config, op_module_fp)

    net = InDuDoNetRatFemur(opt, op_module_fp, op_module_pT).cuda()
    net.load_state_dict(torch.load(os.path.join(config.model_dir, config.name+config.version, 'net_epoch 10.pt')))

    config['len_pre_epoch'] = len(inference_dataloader)
    config.len_pre_epoch = len(inference_dataloader)
    visualizer = Visualizer(config)
    config.name = config.name + "_X"
    visualizer2 = Visualizer(config)  # 这个类用来保存投影域的图

    result_saver = InferenceResultSaver(
        r"..\result\inference_result_csv\RatFemurInDuDoNet",
        config.name
    )

    with torch.no_grad():
        for ii, data in enumerate(inference_dataloader):
            Xgt, Xli, Xma, mask, Sgt, Sma, Sli, Smetal = [x.cuda() for x in data[0:-1]]
            save_path_list = data[-1]
            ListX, ListS, ListYS = net(Xma, Xli, mask, Sma, Sli, Smetal)

            X_final_result = np.clip(ListX[-1].clone().detach().cpu().numpy().squeeze(), 0, 255)
            Xgt_np = Xgt.detach().cpu().numpy().squeeze()

            for j in range(X_final_result.shape[0]):
                save_file = X_final_result[j]
                tmp_gt = Xgt_np[j]

                # mae = computeMAE(save_file / 255., tmp_gt / 255.)
                # ssim = computeSSIM(save_file / 255., tmp_gt / 255.)
                # psnr = computePSNR(save_file / 255., tmp_gt / 255.)
                mae = computeMAE(save_file, tmp_gt)
                ssim = computeSSIM(save_file, tmp_gt)
                psnr = computePSNR(save_file, tmp_gt)
                save_file = save_file.astype(np.uint8)
                cv2.imwrite(save_path_list[j], save_file)
                result_saver.write_one_row([save_path_list[j], mae, ssim, psnr])

            if config['if_show'] and ii % config['show_interval'] == 0:
                show_batch = dict()
                show_batch['Sma_show'] = Sma[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                show_batch['SLI_show'] = Sli[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                S_show_tmp = np.clip(ListYS[-1][0][0].clone().detach().cpu().unsqueeze(0).numpy(), 0, 255)
                # S_show_tmp = (S_show_tmp - np.min(S_show_tmp)) / (np.max(S_show_tmp) - np.min(S_show_tmp)) * 255.0
                show_batch['S_show'] = S_show_tmp
                print(np.max(show_batch['S_show']), np.min(show_batch['S_show']))
                # show_batch['YS_show'] = ListYS[-1][0][0].clone().detach().cpu().unsqueeze(0).numpy().astype(np.uint8)
                show_batch['Sgt'] = Sgt[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                visualizer.reset()
                visualizer.display_current_results(show_batch, 0, ii, True)

                show_batch_2 = dict()
                show_batch_2['Xma_show'] = Xma[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                show_batch_2['XLI_show'] = Xli[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                show_batch_2['X_show'] = np.clip(ListX[-1][0][0].clone().detach().cpu().unsqueeze(0).numpy(), 0, 255)
                print(np.max(show_batch_2['X_show']), np.min(show_batch_2['X_show']))
                show_batch_2['Xgt'] = Xgt[0][0].clone().detach().cpu().unsqueeze(0).numpy()
                visualizer2.reset()
                visualizer2.display_current_results(show_batch_2, 0, ii, True)


if __name__ == '__main__':
    inference()
